import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import csv
import math
df = pd.read_csv('./files/通威股份.csv')[['open','close','low','high','date']]
df['date'] = pd.to_datetime(df['date'])
df['ma5'] = np.nan
df['ma30'] = np.nan

# 设定你想要筛选的日期范围
# start_date = '2021-01-02'
# end_date = '2017-12-15'
end_date = '2024-7-18'
 
# 使用布尔索引进行筛选
# mask = (df['date'] >= start_date) & (df['date'] <= end_date)
mask = df['date'] <= end_date
df = df.loc[mask]

for i in range(4, len(df)):
  df.loc[df.index[i], 'ma5'] = df['close'][i-4:i+1].mean()

for i in range(29, len(df)):
  df.loc[df.index[i], 'ma30'] = df['close'][i-29:i+1].mean()


golden_cross = []
death_cross = []

for i in range(1,len(df)):
  if df['ma5'][i] >= df['ma30'][i] and df['ma5'][i-1] < df['ma30'][i-1]:
    golden_cross.append(df['date'][i].to_pydatetime())
  if df['ma5'][i] <= df['ma30'][i] and df['ma5'][i-1] > df['ma30'][i-1]:
    death_cross.append(df['date'][i].to_pydatetime())

moneys = 200000
shareCertificate = 0

buy_num = 0
sell_num = 0
buy_num_all = 0
sell_num_all = 0

# print(df[df['date'] == golden_cross[0]])

def writeFile(content):
  with open('./files/test.txt', 'a') as file:
    file.write(content)

def sellAndBuy(price, count, money, type, num, date):
  str = f'{type}日期{date}，{type}价格{price}，{type}了{count}手，还剩{money}元，还剩{num}手'
  print(str)
  writeFile(str + '\n')

def calCommission(type, money, num, price):
  if(type == 'sell'):
    allMoney = price * num * 100
    sellCommission = allMoney * 0.00025 + allMoney * 0.00002 + allMoney * 0.001
    if(sellCommission < 5):
      return 5
    else:
      return sellCommission
  else:
    buy_num = math.floor(money / price / 100)
    allMoney = buy_num * price * 100
    sellCommission = allMoney * 0.00025 + allMoney * 0.00002
    if((sellCommission + allMoney) > money):
      buy_num = buy_num - 1
    allMoney = buy_num * price * 100
    sellCommission = allMoney * 0.00025 + allMoney * 0.00002
    if(sellCommission < 5):
      return 5
    else:
      return sellCommission

# print(len(golden_cross))
gold_length = len(golden_cross)
dea_length = len(death_cross)
length = 0
if(gold_length > dea_length):
  length = gold_length
else:
  length = dea_length

for i in range(0, length):
  if(i < dea_length):
    death_cross_price = df[df['date'] == death_cross[i]]['close'].iloc[0]
    if(shareCertificate > 0):
      sell_num = sell_num + 1
      allMoney = shareCertificate * 100 * death_cross_price
      # sellCommission = allMoney * 0.00025 + allMoney * 0.00002 + allMoney * 0.001
      sellCommission = calCommission('sell', moneys, shareCertificate, death_cross_price)
      sell_num_all = sell_num_all + sellCommission
      moneys = moneys + allMoney - sellCommission
      sellAndBuy(death_cross_price, shareCertificate, moneys, '卖出', 0, death_cross[i])
      shareCertificate = 0

  if(i < gold_length):
    golden_cross_price = df[df['date'] == golden_cross[i]]['close'].iloc[0]
    if(moneys > golden_cross_price * 100):
      buy_num = buy_num + 1

      num = math.floor(moneys / golden_cross_price / 100)
      allMoney = num * golden_cross_price * 100
      sellCommission = allMoney * 0.00025 + allMoney * 0.00002
      if((sellCommission + allMoney) > num):
        num = num - 1

      buyMoney = num * 100 * golden_cross_price
      # buyCommission = buyMoney * 0.00025 + buyMoney * 0.00002
      buyCommission = calCommission('buy', moneys, shareCertificate, golden_cross_price)
      buy_num_all = buy_num_all + buyCommission
      shareCertificate = shareCertificate + num
      moneys = moneys - buyMoney - buyCommission
      sellAndBuy(golden_cross_price, num, moneys, '购买', shareCertificate, golden_cross[i])

last_day_golden_cross_price = df[df['date'] == golden_cross[len(golden_cross) - 1]]['close'].iloc[0]
print(last_day_golden_cross_price)

# print(death_cross)
# print(golden_cross)

str = f'''
------------
最终总余额：{moneys}
最终总剩余股份：{shareCertificate * 100}
最终购买次数：{buy_num}
最终总购买手续费：{buy_num_all}
最终卖出次数：{sell_num}
最终总卖出手续费：{sell_num_all}
最后一个金叉股票换钱后总额：{moneys + last_day_golden_cross_price * shareCertificate * 100}
'''
print(str + '\n')
